import torch
import torch.nn as nn

class LSTM(nn.Module):
    def __init__(self, input_dim=1280, hidden_dim=512, out_dim=128, dropout_rate=0.1):
        super(LSTM, self).__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.dropout = nn.Dropout(dropout_rate)
        self.lstm = nn.LSTM(input_size=hidden_dim, hidden_size=hidden_dim, batch_first=True)
        self.fc2 = nn.Linear(hidden_dim, out_dim)
        self.fc3 = nn.Linear(out_dim, out_dim)

    def forward(self, x):
        x = self.fc1(x)
        x = x.unsqueeze(1)
        x, _ = self.lstm(x)
        x = x.squeeze(1)
        x = self.fc2(x)
        x = torch.relu(x)
        x = self.fc3(x)
        x = self.dropout(x)
        x = torch.sigmoid(x)
        return x